#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
性能测试脚本
"""

import time
import pandas as pd
import os
import sys

# 添加项目根目录到Python路径
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from wechat_parser import parse_wechat_bill_xlsx, auto_categorize, batch_auto_categorize
from database import Database

def test_excel_parsing_performance():
    """测试Excel解析性能"""
    print("开始测试Excel解析性能...")
    
    # 使用测试文件
    test_file = '测试/微信支付账单流水文件(20250430-20250731)_20250731143833.xlsx'
    
    if not os.path.exists(test_file):
        print(f"测试文件不存在: {test_file}")
        return
    
    # 测试解析时间
    start_time = time.time()
    transactions = parse_wechat_bill_xlsx(test_file)
    end_time = time.time()
    
    print(f"解析 {len(transactions)} 条交易记录")
    print(f"解析耗时: {end_time - start_time:.2f} 秒")
    
    return transactions

def test_classification_performance(transactions):
    """测试分类性能"""
    print("\n开始测试分类性能...")
    
    # 测试分类时间
    start_time = time.time()
    
    # 提取所有备注
    remarks = [transaction.payee for transaction in transactions]
    
    # 批量分类
    categories = batch_auto_categorize(remarks)
    
    # 将分类结果赋值给交易记录
    for transaction, category in zip(transactions, categories):
        transaction.category = category
        
    end_time = time.time()
    
    print(f"分类 {len(transactions)} 条交易记录")
    print(f"分类耗时: {end_time - start_time:.2f} 秒")

def test_database_performance(transactions):
    """测试数据库性能"""
    print("\n开始测试数据库性能...")
    
    # 初始化数据库
    db = Database('./test_performance.db')
    
    # 测试插入时间
    start_time = time.time()
    for transaction in transactions:
        db.insert_transaction(
            transaction.timestamp,
            transaction.amount,
            transaction.type,
            transaction.category,
            transaction.payee,
            transaction.source,
            transaction.note
        )
    end_time = time.time()
    
    print(f"插入 {len(transactions)} 条交易记录到数据库")
    print(f"插入耗时: {end_time - start_time:.2f} 秒")
    
    # 测试查询时间
    start_time = time.time()
    all_transactions = db.get_transactions()
    end_time = time.time()
    
    print(f"查询 {len(all_transactions)} 条交易记录")
    print(f"查询耗时: {end_time - start_time:.2f} 秒")
    
    # 清理测试数据库
    del db
    os.remove('test_performance.db')

def main():
    print("=== 记账助手性能测试 ===")
    
    # 测试Excel解析性能
    transactions = test_excel_parsing_performance()
    
    if transactions:
        # 测试分类性能
        test_classification_performance(transactions)
        
        # 测试数据库性能
        test_database_performance(transactions)
    
    print("\n=== 性能测试完成 ===")

if __name__ == "__main__":
    main()